
import torch as th
from torch import nn
from dgl.utils import expand_as_pair
from dgl.nn.pytorch.utils import Identity

from .nn.functional import edge_softmax, gat_propagate, gat_apply_edge


class GATConv(th.nn.Module):

    def __init__(self,
                 in_feats,
                 out_feats,
                 num_heads,
                 feat_drop=0.,
                 attn_drop=0.,
                 negative_slope=0.2,
                 residual=False,
                 activation=None,
                 bias=True):
        super().__init__()
        self._num_heads = num_heads
        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
        self._out_feats = out_feats
        if isinstance(in_feats, tuple):
            self.fc_src = nn.Linear(
                self._in_src_feats, out_feats * num_heads, bias=False)
            self.fc_dst = nn.Linear(
                self._in_dst_feats, out_feats * num_heads, bias=False)
        else:
            self.fc = nn.Linear(
                self._in_src_feats, out_feats * num_heads, bias=False)

        self.attn_l = nn.Parameter(th.FloatTensor(size=(1, num_heads, out_feats)))
        self.attn_r = nn.Parameter(th.FloatTensor(size=(1, num_heads, out_feats)))
        self.feat_drop = nn.Dropout(feat_drop)
        self.attn_drop = nn.Dropout(attn_drop)
        self.leaky_relu = nn.LeakyReLU(negative_slope)
        if bias:
            self.bias = nn.Parameter(th.FloatTensor(size=(num_heads * out_feats,)))
        else:
            self.register_buffer('bias', None)
        if residual:
            if self._in_dst_feats != out_feats * num_heads:
                self.res_fc = nn.Linear(
                    self._in_dst_feats, num_heads * out_feats, bias=False)
            else:
                self.res_fc = Identity()
        else:
            self.register_buffer('res_fc', None)
        self.reset_parameters()
        self.activation = activation


    def reset_parameters(self):
        gain = nn.init.calculate_gain('relu')
        if hasattr(self, 'fc'):
            nn.init.xavier_normal_(self.fc.weight, gain=gain)
        else:
            nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
            nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
        nn.init.xavier_normal_(self.attn_l, gain=gain)
        nn.init.xavier_normal_(self.attn_r, gain=gain)
        if self.bias is not None:
            nn.init.constant_(self.bias, 0)
        if isinstance(self.res_fc, nn.Linear):
            nn.init.xavier_normal_(self.res_fc.weight, gain=gain)


    def forward(self, graph, feat):
        if isinstance(feat, tuple):
            src_prefix_shape = feat[0].shape[:-1]
            dst_prefix_shape = feat[1].shape[:-1]
            h_src = self.feat_drop(feat[0])
            h_dst = self.feat_drop(feat[1])
            if not hasattr(self, 'fc_src'):
                feat_src = self.fc(h_src).view(
                    *src_prefix_shape, self._num_heads, self._out_feats)
                feat_dst = self.fc(h_dst).view(
                    *dst_prefix_shape, self._num_heads, self._out_feats)
            else:
                feat_src = self.fc_src(h_src).view(
                    *src_prefix_shape, self._num_heads, self._out_feats)
                feat_dst = self.fc_dst(h_dst).view(
                    *dst_prefix_shape, self._num_heads, self._out_feats)
        else:
            src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
            h_src = h_dst = self.feat_drop(feat)
            feat_src = feat_dst = self.fc(h_src).view(
                *src_prefix_shape, self._num_heads, self._out_feats)

        el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(dim=-1)
        er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(dim=-1)
        # split the attention
        rsts = []
        for head_i in range(self._num_heads):
            eu = el[:,head_i,:].clone()
            ev = er[:,head_i,:].clone()
            att_raw = self.leaky_relu(gat_apply_edge(graph, None, eu, ev))
            att = self.attn_drop(edge_softmax(graph, att_raw))
            rsts.append(gat_propagate(graph, att, feat_src[:,head_i,:]))

        rst = th.stack(rsts, dim=1).reshape(feat_src.shape)

        # residual
        if self.res_fc is not None:
            # Use -1 rather than self._num_heads to handle broadcasting
            resval = self.res_fc(h_dst).view(*dst_prefix_shape, -1, self._out_feats)
            rst = rst + resval
        # bias
        if self.bias is not None:
            rst = rst + self.bias.view(
                *((1,) * len(dst_prefix_shape)), self._num_heads, self._out_feats)
        # activation
        if self.activation:
            rst = self.activation(rst)

        return rst